package com.unicdata.twofactorauth.utils;

import com.google.zxing.BarcodeFormat;
import com.google.zxing.client.j2se.MatrixToImageWriter;
import com.google.zxing.common.BitMatrix;
import com.google.zxing.qrcode.QRCodeWriter;
import com.unicdata.twofactorauth.config.TwoFactorProperties;
import com.unicdata.twofactorauth.model.CompleteTwoFactorSetupResult;
import com.unicdata.twofactorauth.model.RecoveryCodesResult;
import com.unicdata.twofactorauth.model.TwoFactorSecretResult;
import com.unicdata.twofactorauth.model.TwoFactorSetupResult;
import com.unicdata.twofactorauth.storage.TwoFactorStorage;
import com.warrenstrange.googleauth.GoogleAuthenticator;
import com.warrenstrange.googleauth.GoogleAuthenticatorConfig;
import com.warrenstrange.googleauth.KeyRepresentation;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileSystems;
import java.nio.file.Path;
import java.security.SecureRandom;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * 双因子认证工具类
 *
 * @author lhn
 * @date 2025/07/02 13:46
 * @since 1.0.0
 **/
@Slf4j
public class TwoFactorAuthUtils {

    private final GoogleAuthenticator googleAuthenticator;
    /**
     * -- GETTER --
     *  获取当前配置
     */
    @Getter
    private final TwoFactorProperties properties;
    /**
     * -- GETTER --
     *  获取存储实例
     */
    @Getter
    private final TwoFactorStorage storage;
    private final SecureRandom secureRandom;

    /**
     * 使用默认配置构造（不带存储）
     */
    public TwoFactorAuthUtils() {
        this(new TwoFactorProperties(), null);
    }

    /**
     * 使用自定义配置构造（不带存储）
     */
    public TwoFactorAuthUtils(TwoFactorProperties properties) {
        this(properties, null);
    }

    /**
     * 使用自定义配置和存储构造
     */
    public TwoFactorAuthUtils(TwoFactorProperties properties, TwoFactorStorage storage) {
        this.properties = properties;
        this.storage = storage;
        this.googleAuthenticator = createGoogleAuthenticator(properties);
        this.secureRandom = new SecureRandom();
    }

    /**
     * 创建GoogleAuthenticator实例
     */
    private GoogleAuthenticator createGoogleAuthenticator(TwoFactorProperties properties) {
        GoogleAuthenticatorConfig config = new GoogleAuthenticatorConfig.GoogleAuthenticatorConfigBuilder()
                .setTimeStepSizeInMillis(TimeUnit.SECONDS.toMillis(properties.getTimeStepSizeInSeconds()))
                .setWindowSize(properties.getWindowSize())
                .setCodeDigits(properties.getCodeDigits())
                .setKeyRepresentation(getKeyRepresentation(properties.getKeyRepresentation()))
                .build();
        return new GoogleAuthenticator(config);
    }

    /**
     * 获取密钥表示方式枚举
     */
    private KeyRepresentation getKeyRepresentation(String keyRepresentation) {
        try {
            return KeyRepresentation.valueOf(keyRepresentation.toUpperCase());
        } catch (IllegalArgumentException e) {
            log.warn("无效的密钥表示方式: {}, 使用默认值BASE32", keyRepresentation);
            return KeyRepresentation.BASE32;
        }
    }

    /**
     * 生成2FA密钥和二维码URL
     */
    public TwoFactorSecretResult generateTwoFactorSecret(String username) {
        // 生成密钥
        final String secretKey = googleAuthenticator.createCredentials().getKey();

        try {
            String encodedUsername = URLEncoder.encode(username, StandardCharsets.UTF_8.toString());
            // 生成二维码URL
            String otpAuthURL = String.format(
                    "otpauth://totp/%s?secret=%s&issuer=%s",
                    encodedUsername,
                    secretKey,
                    URLEncoder.encode(properties.getIssuer(), StandardCharsets.UTF_8.toString())
            );
            return new TwoFactorSecretResult(secretKey, otpAuthURL);
        } catch (UnsupportedEncodingException e) {
            throw new RuntimeException("生成2FA密钥失败", e);
        }
    }

    /**
     * 生成恢复码
     *
     * @param count 恢复码数量（默认8个）
     * @return 恢复码结果
     */
    public RecoveryCodesResult generateRecoveryCodes(int count) {
        List<String> codes = new ArrayList <>();
        for (int i = 0; i < count; i++) {
            codes.add(generateSingleRecoveryCode());
        }
        return new RecoveryCodesResult(codes, System.currentTimeMillis());
    }

    /**
     * 生成恢复码（使用默认数量）
     */
    public RecoveryCodesResult generateRecoveryCodes() {
        return generateRecoveryCodes(8);
    }

    /**
     * 生成单个恢复码
     */
    private String generateSingleRecoveryCode() {
        // 生成8位随机字符串，包含数字和字母
        String chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < 8; i++) {
            sb.append(chars.charAt(secureRandom.nextInt(chars.length())));
        }
        return sb.toString();
    }

    /**
     * 生成二维码图片（Base64编码）
     *
     * @param otpAuthURL OTP认证URL
     * @param width 二维码宽度（默认200）
     * @param height 二维码高度（默认200）
     * @return Base64编码的PNG图片
     */
    public String generateQRCodeImage(String otpAuthURL, int width, int height) {
        try {
            QRCodeWriter qrCodeWriter = new QRCodeWriter();
            BitMatrix bitMatrix = qrCodeWriter.encode(otpAuthURL, BarcodeFormat.QR_CODE, width, height);

            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            MatrixToImageWriter.writeToStream(bitMatrix, "PNG", outputStream);

            byte[] qrCodeBytes = outputStream.toByteArray();
            return Base64.getEncoder().encodeToString(qrCodeBytes);
        } catch (Exception e) {
            log.error("生成二维码失败", e);
            throw new RuntimeException("生成二维码失败", e);
        }
    }

    /**
     * 生成二维码图片（Base64编码）- 使用配置文件中的尺寸
     */
    public String generateQRCodeImage(String otpAuthURL) {
        return generateQRCodeImage(otpAuthURL, properties.getQrCode().getWidth(), properties.getQrCode().getHeight());
    }

    /**
     * 生成二维码图片文件
     *
     * @param otpAuthURL OTP认证URL
     * @param filePath 文件保存路径
     * @param width 二维码宽度
     * @param height 二维码高度
     */
    public void generateQRCodeImageFile(String otpAuthURL, String filePath, int width, int height) {
        try {
            QRCodeWriter qrCodeWriter = new QRCodeWriter();
            BitMatrix bitMatrix = qrCodeWriter.encode(otpAuthURL, BarcodeFormat.QR_CODE, width, height);

            Path path = FileSystems.getDefault().getPath(filePath);
            MatrixToImageWriter.writeToPath(bitMatrix, "PNG", path);
        } catch (Exception e) {
            log.error("生成二维码文件失败", e);
            throw new RuntimeException("生成二维码文件失败", e);
        }
    }

    /**
     * 生成二维码图片文件 - 使用配置文件中的尺寸
     */
    public void generateQRCodeImageFile(String otpAuthURL, String filePath) {
        generateQRCodeImageFile(otpAuthURL, filePath, properties.getQrCode().getWidth(), properties.getQrCode().getHeight());
    }

    /**
     * 生成完整的2FA设置信息（包含密钥、URL和二维码）
     *
     * @param username 用户名
     * @param qrCodeWidth 二维码宽度
     * @param qrCodeHeight 二维码高度
     * @return 2FA设置结果
     */
    public TwoFactorSetupResult generateTwoFactorSetup(String username, int qrCodeWidth, int qrCodeHeight) {
        TwoFactorSecretResult secretResult = generateTwoFactorSecret(username);
        String qrCodeBase64 = generateQRCodeImage(secretResult.getOtpAuthUrl(), qrCodeWidth, qrCodeHeight);

        return new TwoFactorSetupResult(
                secretResult.getSecret(),
                secretResult.getOtpAuthUrl(),
                qrCodeBase64
        );
    }

    /**
     * 生成完整的2FA设置信息 - 使用配置文件中的二维码尺寸
     */
    public TwoFactorSetupResult generateTwoFactorSetup(String username) {
        return generateTwoFactorSetup(username, properties.getQrCode().getWidth(), properties.getQrCode().getHeight());
    }

    /**
     * 生成完整的2FA设置信息（包含恢复码）
     *
     * @param username 用户名
     * @param qrCodeWidth 二维码宽度
     * @param qrCodeHeight 二维码高度
     * @param recoveryCodeCount 恢复码数量
     * @return 完整2FA设置结果
     */
    public CompleteTwoFactorSetupResult generateCompleteTwoFactorSetup(String username, int qrCodeWidth, int qrCodeHeight, int recoveryCodeCount) {
        TwoFactorSetupResult setupResult = generateTwoFactorSetup(username, qrCodeWidth, qrCodeHeight);
        RecoveryCodesResult recoveryResult = generateRecoveryCodes(recoveryCodeCount);

        return new CompleteTwoFactorSetupResult(
                setupResult.getSecret(),
                setupResult.getOtpAuthUrl(),
                setupResult.getQrCodeBase64(),
                recoveryResult.getCodes(),
                recoveryResult.getTimestamp()
        );
    }

    /**
     * 生成完整的2FA设置信息（包含恢复码）- 使用默认参数
     */
    public CompleteTwoFactorSetupResult generateCompleteTwoFactorSetup(String username) {
        return generateCompleteTwoFactorSetup(username, properties.getQrCode().getWidth(), properties.getQrCode().getHeight(), 8);
    }

    /**
     * 验证2FA码
     */
    public boolean validateCode(String secretKey, int code) {
        try {
            return googleAuthenticator.authorize(secretKey, code);
        } catch (Exception e) {
            log.error("2FA验证失败", e);
            return false;
        }
    }

    /**
     * 验证2FA码（字符串形式）
     */
    public boolean validateCode(String secretKey, String code) {
        try {
            int intCode = Integer.parseInt(code);
            return validateCode(secretKey, intCode);
        } catch (NumberFormatException e) {
            log.error("验证码格式错误: {}", code);
            return false;
        }
    }

    /**
     * 验证恢复码
     */
    public boolean validateRecoveryCode(List<String> recoveryCodes, String inputCode) {
        if (recoveryCodes == null || inputCode == null) {
            return false;
        }
        return recoveryCodes.contains(inputCode.toUpperCase());
    }

    // ============ 存储相关方法（需要存储实现） ============

    /**
     * 保存2FA设置到存储（需要存储实现）
     */
    public boolean saveTwoFactorSetup(String userId, String secret, List<String> recoveryCodes) {
        if (storage == null) {
            throw new IllegalStateException("未配置存储实现，无法保存2FA设置");
        }
        boolean secretSaved = storage.saveSecret(userId, secret);
        boolean codesSaved = storage.saveRecoveryCodes(userId, recoveryCodes);
        return secretSaved && codesSaved;
    }

    /**
     * 验证用户2FA码（从存储获取密钥）
     */
    public boolean validateUserCode(String userId, String code) {
        if (storage == null) {
            throw new IllegalStateException("未配置存储实现，无法验证用户2FA码");
        }
        Optional <String> secret = storage.getSecret(userId);
        if (!secret.isPresent()) {
            return false;
        }
        return validateCode(secret.get(), code);
    }

    /**
     * 验证用户恢复码（从存储获取并使用）
     */
    public boolean validateUserRecoveryCode(String userId, String recoveryCode) {
        if (storage == null) {
            throw new IllegalStateException("未配置存储实现，无法验证用户恢复码");
        }
        return storage.useRecoveryCode(userId, recoveryCode);
    }

    /**
     * 禁用用户2FA
     */
    public boolean disableTwoFactor(String userId) {
        if (storage == null) {
            throw new IllegalStateException("未配置存储实现，无法禁用2FA");
        }
        boolean secretDeleted = storage.deleteSecret(userId);
        boolean codesDeleted = storage.deleteAllRecoveryCodes(userId);
        boolean statusUpdated = storage.setTwoFactorEnabled(userId, false);
        return secretDeleted && codesDeleted && statusUpdated;
    }

    /**
     * 重新生成恢复码
     */
    public RecoveryCodesResult regenerateRecoveryCodes(String userId, int count) {
        if (storage == null) {
            throw new IllegalStateException("未配置存储实现，无法重新生成恢复码");
        }
        RecoveryCodesResult result = generateRecoveryCodes(count);
        storage.saveRecoveryCodes(userId, result.getCodes());
        return result;
    }

    /**
     * 静态工厂方法 - 使用默认配置创建实例
     */
    public static TwoFactorAuthUtils createDefault() {
        return new TwoFactorAuthUtils();
    }

    /**
     * 静态工厂方法 - 使用自定义配置创建实例
     */
    public static TwoFactorAuthUtils create(TwoFactorProperties properties) {
        return new TwoFactorAuthUtils(properties);
    }

    /**
     * 静态工厂方法 - 使用自定义配置和存储创建实例
     */
    public static TwoFactorAuthUtils create(TwoFactorProperties properties, TwoFactorStorage storage) {
        return new TwoFactorAuthUtils(properties, storage);
    }
}